#include "preprocessor.h"

#include <filesystem>

#include "dwarf_utils.h"
#include "elog.h"
#include "elog_file.h"
#include "live_stacking/star_stack.h"

using namespace cv;

int Preprocessor::readLightFrame(const cv::Mat& img_light, cv::Mat& img_color) {
    if (!img_light.empty()) {
        demosaic(img_light, img_color);
        return 0;
    } else {
        return -1;
    }
}

int Preprocessor::readLightFrame(const char* img_path, cv::Mat& img_color) {
    cv::Mat img_light;
    ushort* img_light_data = nullptr;
    file_manager_.convertTo16BitImg(img_path, img_light_data, img_light);

    if (!img_light.empty()) {
        demosaic(img_light, img_color);
        xfree(img_light_data);

        return 0;
    } else {
        return -1;
    }
}

int Preprocessor::calibrateLightFrame(const char* img_path, const cv::Mat img_bad_pix_map, cv::Mat& img_color) {
    cv::Mat img_light;
    ushort* img_light_data = nullptr;
    file_manager_.convertTo16BitImg(img_path, img_light_data, img_light);

    if (!img_light.empty()) {
        badPixCorrection(img_light, img_bad_pix_map);
        demosaic(img_light, img_color);
        unlinkedStretch(img_color);

        xfree(img_light_data);

        return 0;
    } else {
        return -1;
    }
}

int Preprocessor::readBadPixMap(const std::string& dark_folder, int user_set_gain, int binning, cv::Mat& img_bad_pix_map, double user_set_exp) {
    std::ostringstream img_bad_pix_map_path;
    img_bad_pix_map_path << dark_folder << "badPix_exp_" << std::fixed << std::setprecision(6) << user_set_exp << "_gain_" << user_set_gain << "_bin_" << binning + 1 << "_thre_0.5.png";

    ushort* img_bad_pix_map_data = nullptr;
    file_manager_.convertTo16BitImg(img_bad_pix_map_path.str().c_str(), img_bad_pix_map_data, img_bad_pix_map);

    if (img_bad_pix_map.empty()) {
        std::ostringstream img_dark_path;
        img_dark_path << dark_folder << "exp_" << std::fixed << std::setprecision(6) << user_set_exp << "_gain_" << user_set_gain << "_bin_" << binning + 1 << ".png";

        cv::Mat img_dark;
        ushort* img_dark_data = nullptr;
        file_manager_.convertTo16BitImg(img_dark_path.str().c_str(), img_dark_data, img_dark);

        if (img_dark.empty()) {
            return -1;
        } else {
            log_i("Generate bad pixel map");
            generateBadPixMap(img_dark, img_bad_pix_map);
            imwrite(img_bad_pix_map_path.str(), img_bad_pix_map);
            system("sync");
        }
    }

    return 0;
}

int Preprocessor::generateBadPixMap(const Mat& imgDark, Mat& imgBadPix) {
	if(imgDark.empty())
		return -1;

	Mat matMean, matStddev;
	meanStdDev(imgDark, matMean, matStddev);
	double mean = matMean.at<double>(0,0);
	double stddev = matStddev.at<double>(0,0);
	double threshold = mean + 0.5 * stddev;
	log_i("mean:%f", mean);

	cv::threshold(imgDark, imgBadPix, threshold, 65535, THRESH_TOZERO);
	return 0;
}

int Preprocessor::badPixCorrection(Mat& imgSrc, const Mat& imgBadPix) {
	Mat imgSrcBorder;
	copyMakeBorder(imgSrc, imgSrcBorder, 2, 2, 2, 2, cv::BORDER_WRAP);		// 按滚动方式扩展边界
    
	if(!imgBadPix.empty()) {												// 暗场去坏点
		imgBadPix.forEach<ushort>([&](ushort& pixel, const int position[]) {
			if(pixel) {
				int y = position[0];
				int x = position[1];

				ushort* imgSrcRowPtr = imgSrc.ptr<ushort>(y);
				imgSrcRowPtr[x] = bayerMedian(imgSrcBorder, Point(x + 2, y + 2));
			}
		});
	}

	std::vector<cv::Point> pixelDefects;									// 动态去坏点
	badPixDynamicDetection(imgSrc, pixelDefects);

	std::for_each (pixelDefects.begin(), pixelDefects.end(), [&](const cv::Point& point) {
		imgSrc.at<ushort>(point) = bayerMedian(imgSrcBorder, Point(point.x + 2, point.y + 2));
	});

	return 0;
}

void Preprocessor::calculateResult(const cv::Point& point, const cv::Mat& image, std::vector<cv::Point>& defects) {
    int x = point.x;
    int y = point.y;
    int cols = image.cols;
    int rows = image.rows;

    // 检查像素点是否在图像范围内
    if (x >= 1 && x < cols - 1 && y >= 1 && y < rows - 1) {
        // 获取上下左右的像素点位置
        cv::Point up(x, y - 1);
        cv::Point down(x, y + 1);
        cv::Point left(x - 1, y);
        cv::Point right(x + 1, y);

        // 计算上下左右像素的值
        ushort upValue = image.at<ushort>(up);
        ushort downValue = image.at<ushort>(down);
        ushort leftValue = image.at<ushort>(left);
        ushort rightValue = image.at<ushort>(right);
        ushort centerValue = image.at<ushort>(point);

        // 计算结果并存储
        int result = centerValue * 4 - (upValue + downValue + leftValue + rightValue);
		if(result > 9000) {
			defects.push_back(point);
		}
    }
}

void Preprocessor::locatePixelDefects(cv::Mat& binaryImage, std::vector<cv::Point>& defects) {
    cv::Mat labels, stats, centroids;
    int numLabels = cv::connectedComponentsWithStats(binaryImage, labels, stats, centroids);

    for (int label = 1; label < numLabels; ++label) {
        int* statsPtr = stats.ptr<int>(label);
        int area = statsPtr[cv::ConnectedComponentsTypes::CC_STAT_AREA];

        // 孤立像素坏点
        if (area == 1) {
            double* centroidPtr = centroids.ptr<double>(label);
            int x = static_cast<int>(centroidPtr[0]);
            int y = static_cast<int>(centroidPtr[1]);
            defects.emplace_back(x, y);
        }
    }
}

void Preprocessor::badPixDynamicDetection(const Mat& src, std::vector<cv::Point>& truePixelDefects) {
	Mat matMean, matStddev;
	meanStdDev(src, matMean, matStddev);
	double mean = matMean.at<double>(0,0);
	double stddev = matStddev.at<double>(0,0);
	double threshold = (mean + 3.0 * stddev) > 65534 ? 65534 : (mean + 3.0 * stddev);

	cv::Mat binaryImage;
    cv::threshold(src, binaryImage, threshold, 65535, cv::THRESH_BINARY);

	Mat kernel = getStructuringElement(MORPH_ELLIPSE, Size(5, 5), Point(-1, -1));
	morphologyEx(binaryImage, binaryImage, MORPH_CLOSE, kernel);

	cv::Mat binaryImage8bit;
	binaryImage.convertTo(binaryImage8bit, CV_8U, 1.0 / 256);

    // 定位像素坏点
    std::vector<cv::Point> pixelDefects;
    locatePixelDefects(binaryImage8bit, pixelDefects);

	std::for_each (pixelDefects.begin(), pixelDefects.end(), [&](const cv::Point& point) {
		calculateResult(point, src, truePixelDefects);
	});
}

ushort Preprocessor::bayerMedian(const Mat& input, const cv::Point& center) {
	int x = center.x;
	int y = center.y;

	std::vector<ushort> values;
	for (int dy = -2; dy <= 2; dy++) {					// 取bayer阵列5*5同色像素
		const ushort* inputRowPtr = input.ptr<ushort>(y + dy);

		for (int dx = -2; dx <= 2; dx++) {
			if ((dx == -1 && dy == -2) || (dx == 1 && dy == -2) || 
				(dy == -1) || 
				(dx == -1 && dy == 0) || (dx == 0 && dy == 0) || (dx == 1 && dy == 0) || 
				(dy == 1) ||  
				(dx == -1 && dy == 2) || (dx == 1 && dy == 2))
				continue;
			
			values.push_back(inputRowPtr[x + dx]);
		}
	}
	std::sort(values.begin(), values.end());
	return (values[3] + values[4]) / 2;
}

int Preprocessor::demosaic(const Mat& img, Mat& img_color) {
	if(img.empty())
		return -1;
	if(img.channels() != 1)
		return -2;

	Mat mv[3];
	cvtColor(img, img_color, COLOR_BayerGB2RGB);		// 去马赛克
	normalize(img_color, img_color, 0, 65535, cv::NORM_MINMAX);
	split(img_color, mv);
	mv[0] /= 0.377734;
	mv[2] /= 0.619960;
	merge(mv, 3, img_color);

	return 0;
}
